/**
 * Copyright 2012 Brigham Young University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package edu.byu.nlp.cluster.em;

import java.util.logging.Logger;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;

import edu.byu.nlp.cluster.Dataset;
import edu.byu.nlp.math.optimize.IterativeOptimizer;
import edu.byu.nlp.math.optimize.IterativeOptimizer.Optimizable;
import edu.byu.nlp.math.optimize.IterativeOptimizer.ReturnType;
import edu.byu.nlp.math.optimize.ValueAndObject;

/**
 * Class for performing expectation maximization on a mixture model. Uses the visitor pattern via the
 * {@code PosteriorVisitor} interface for efficiently handling a potentially large number of posteriors.
 * 
 * @author rah67
 *
 */
public class ExpectationMaximization {

	private static final Logger logger = Logger.getLogger(ExpectationMaximization.class.getName());
	
	private final IterativeOptimizer optimizer;
	
	public ExpectationMaximization(IterativeOptimizer optimizer) {
		this.optimizer = optimizer;
	}
	
	private static final class EMOptimizable<P> implements Optimizable<P> {

		private final Dataset data;
		private final Expectable<P> expectable;
		// This variable is used to "alternate" storage locations of the parameters in order to minimize memory usage.
		private ParameterPair<P> paramPair;
		
		public EMOptimizable(Dataset data, Expectable<P> expectable, ParameterPair<P> paramPair) {
			this.data = data;
			this.expectable = expectable;
			this.paramPair = paramPair;
		}

		/** {@inheritDoc} */
		@Override
		public ValueAndObject<P> computeNext(P curParams) {
			paramPair.setCurParams(curParams);
			ValueAndObject<P> vao = emSingleIteration(data, expectable, paramPair.getCurParams(), paramPair.getNextParams());
			paramPair.swap();
			paramPair.setCurParams(vao.getObject());
			return vao;
		}
		
	}
	
	public static class ParameterPair<P> {
		
		private P curParams;
		private P nextParams;
		
		public ParameterPair(P curParams, P nextParams) {
			Preconditions.checkNotNull(curParams);

			this.curParams = curParams;
			this.nextParams = nextParams;
		}

		public void swap() {
			P tmp = curParams;
			curParams = nextParams;
			nextParams = tmp;
		}
		
		public P getCurParams() {
			return curParams;
		}
		
		public void setCurParams(P curParams) {
			this.curParams = curParams;
		}
		
		public P getNextParams() {
			return nextParams;
		}
		
		public void setNextParams(P nextParams) {
			this.nextParams = nextParams;
		}
	}
	
	/**
	 * Performs expectation maximization over the specified dataset using the given expectable and initial parameters.
	 * When this method returns, paramPair.getNextParams() will be equal to (but not necessarily the same as) the object
	 * in the returned value, object pair. 
	 */
	public <P> ValueAndObject<P> em(Dataset data, Expectable<P> expectable, ParameterPair<P> paramPair) {
		Preconditions.checkNotNull(data);
		Preconditions.checkNotNull(expectable);
		Preconditions.checkNotNull(paramPair);
		
		EMOptimizable<P> optimizable = new EMOptimizable<P>(data, expectable, paramPair);
		ValueAndObject<P> vao = optimizer.optimize(optimizable, ReturnType.LAST, true, paramPair.getCurParams());
		
		// optimizable.computeNext() leaves the parameters swapped, so swap them one last time.
		paramPair.swap();

		return vao;
	}
	
	/**
	 * Performs a single iteration of expectation maximization.
	 */
	@VisibleForTesting
	static <P> ValueAndObject<P> emSingleIteration(Dataset data, Expectable<P> expectable, P curParams,
			P nextParams) {
		// E-step
		Maximizable<P> maximizable = expectable.expect(data, curParams, nextParams);
		// M-step
		return maximizable.maximize();
	}

	/**
	 * Performs annealed expectation maximization over the specified dataset using the given annealable and initial
	 * parameters. When this method returns, paramPair.getNextParams() will be equal to (but not necessarily the same
	 * as) the object in the returned value, object pair. 
	 */
	public <P> ValueAndObject<P> annealedEm(Dataset data, final Annealable<P> annealable, ParameterPair<P> paramPair,
			double... annealingSchedule) {
		Preconditions.checkNotNull(data);
		Preconditions.checkNotNull(annealable);
		Preconditions.checkNotNull(paramPair);
		Preconditions.checkNotNull(annealingSchedule);
		
		ValueAndObject<P> parametersAndLowerbound = null;
		for (double temp : annealingSchedule) {
			logger.info("Setting temperature = " + temp);
			Expectable<P> expectable = annealable.annealedExpectable(temp);
			parametersAndLowerbound = em(data, expectable, paramPair);
			logger.info(String.format("Lower bound for temperature %f = %f", temp, parametersAndLowerbound.getValue()));
			
			// Allow for less memory usage by potentially swapping between two sets of parameters
			if (!paramPair.getNextParams().equals(parametersAndLowerbound.getObject())) {
				throw new IllegalStateException("parameters are NOT the same");
			}
			paramPair.swap();
			// Ensure that the curParams are the value that is actually returned (in the case that the implementer is
			// not reusing).
			paramPair.setCurParams(parametersAndLowerbound.getObject());
		}
		
		// The last two lines of the for-loop put the cur/next pairs in the inverse order that we need them.
		if (annealingSchedule.length > 0) {
			paramPair.swap();
		}
		return parametersAndLowerbound;
	}
	
}
